{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.document_loaders import PyPDFLoader\n",
    "from typing import List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! pip install -U text2vec\n",
    "#! pip install modelscope\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#模型下载\n",
    "#from modelscope import snapshot_download\n",
    "#model_dir = snapshot_download('Jerry0/text2vec-base-chinese')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m2024-07-11 09:41:24.278\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mtext2vec.sentence_model\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m80\u001b[0m - \u001b[34m\u001b[1mUse device: cpu\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-4.4401505e-04 -2.9734740e-01  8.5790145e-01 ... -5.2770108e-01\n",
      "  -1.4315575e-01 -1.0007919e-01]\n",
      " [ 6.5362024e-01 -7.6666206e-02  9.5962244e-01 ... -6.0122508e-01\n",
      "  -1.6802035e-03  2.1457647e-01]]\n"
     ]
    }
   ],
   "source": [
    "from text2vec import SentenceModel\n",
    "sentences = ['如何更换花呗绑定银行卡', '花呗更改绑定银行卡']\n",
    "\n",
    "embeddings_model = SentenceModel('D:/06git/01python/02vanna/03Chroma/embedding/Jerry0/text2vec-base-chinese')\n",
    "embeddings = embeddings_model.encode(sentences)\n",
    "print(embeddings)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "chroma_client = chromadb.PersistentClient(path=\"./db\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = PyPDFLoader(\"D:/05work/01telchina/04智能灯杆/005灯杆2.0/泰华智慧灯杆V2.0系统用户使用手册.pdf\")\n",
    "pages = loader.load_and_split()\n",
    "texts = [p.page_content for p in pages]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ids=[\"id_{}\".format(i) for i in range(len(texts))] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = embeddings_model.encode(texts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(77, 768)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeddings.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "collection = chroma_client.get_or_create_collection(name=\"smartpole\")\n",
    "\n",
    "collection.add(\n",
    "    documents=texts,\n",
    "    embeddings=embeddings_model.encode(texts),\n",
    "    ids=[\"id_{}\".format(i) for i in range(len(texts))] \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! pip install langchain_chroma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "langchain_chroma = Chroma(\n",
    "    client=chroma_client,\n",
    "    collection_name=\"smartpole\",\n",
    "    embedding_function=embeddings_model.get_sentence_embeddings,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'smartpole'"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "langchain_chroma._collection_name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prompt(question, context):\n",
    "    prompt = \"\"\" 基于以下已知的信息, 专业、简要的回答用户的问题,\n",
    "            如果无法从提供的内容中获取答案, 请说: \"知识库中提供的内容不足以回答此问题\" 禁止胡乱编造, 回答的时候最好按照1.2.3.点进行总结。 \n",
    "            已知内容: \n",
    "            {context}\n",
    "            问题:\n",
    "            {question},请使用和用户相同的语言进行回答.\n",
    "        \"\"\"\n",
    "    return prompt.format(question=question, context=context)\n",
    "\n",
    "def get_prompt2(question, context):\n",
    "    prompt = \"\"\"\n",
    "        已知信息：\n",
    "        {context}\n",
    "        回答要求：\n",
    "        - 请使用简洁且专业的语言来回答用户的问题。\n",
    "        - 如果你不知道答案，请回答“没有在知识库中查找到相关信息，建议咨询相关技术支持或参考官方文档进行操作”。\n",
    "        - 避免提及你是从已知信息中获得的知识。\n",
    "        - 请保证答案与已知信息中描述的一致。\n",
    "        - 请使用 Markdown 语法优化答案的格式。\n",
    "        - 已知信息中的图片、链接地址和脚本语言请直接返回。\n",
    "        - 请使用与问题相同的语言来回答。\n",
    "        问题：\n",
    "        {question}\n",
    "    \"\"\"\n",
    "    return prompt.format(question=question, context=context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jwt\n",
    "import time\n",
    "from langchain_openai import ChatOpenAI\n",
    "zhipuai_api_key = \"4b60409403162df4e015cdd99af1d803.aQEEaBQBpY6AXqPL\"\n",
    "\n",
    "def generate_token(apikey: str, exp_seconds: int):\n",
    "        try:\n",
    "            id, secret = apikey.split(\".\")\n",
    "        except Exception as e:\n",
    "            raise Exception(\"invalid apikey\", e)\n",
    "\n",
    "        payload = {\n",
    "            \"api_key\": id,\n",
    "            \"exp\": int(round(time.time() * 1000)) + exp_seconds * 1000,\n",
    "            \"timestamp\": int(round(time.time() * 1000)),\n",
    "        }\n",
    "\n",
    "        return jwt.encode(\n",
    "            payload,\n",
    "            secret,\n",
    "            algorithm=\"HS256\",\n",
    "            headers={\"alg\": \"HS256\", \"sign_type\": \"SIGN\"},\n",
    "        )\n",
    "\n",
    "zhipu_llm = ChatOpenAI(\n",
    "        model_name=\"glm-4\",\n",
    "        openai_api_base=\"https://open.bigmodel.cn/api/paas/v4\",\n",
    "        #openai_api_base=\"http://localhost:1234/v1\",\n",
    "        openai_api_key=generate_token(zhipuai_api_key,864000),\n",
    "        streaming=False,\n",
    "        verbose=True\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. 打开灯杆管理页面。\n",
      "2. 点击新增按钮，这将打开新增灯杆页面。\n",
      "3. 在此页面上填写相应的信息，包括但不限于灯杆类型名称、产品名称、产品类型、产品型号和说明。\n",
      "4. 填写完毕后，点击保存按钮，即可保存新录入的灯杆信息。 \n",
      "\n",
      "请注意按照用户使用手册中的指导步骤进行操作。\n"
     ]
    }
   ],
   "source": [
    "question = \"添加灯杆?\"\n",
    "query_embedding = embeddings_model.encode([question])\n",
    "documents = langchain_chroma.similarity_search_by_vector_with_relevance_scores(query_embedding,k=7)\n",
    "content = '';\n",
    "for doc in documents:\n",
    "    content +=doc[0].page_content\n",
    "\n",
    "messages = get_prompt(context=content,question=question)\n",
    "response = zhipu_llm.invoke(messages)\n",
    "print(response.content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
